{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "respiratory-retirement",
   "metadata": {},
   "source": [
    "# WORK IN PROGRESS\n",
    "\n",
    "Use Beautify.ipynb instead"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "thousand-pocket",
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import numpy as np\n",
    "import torch\n",
    "from PIL import Image, ImageDraw, ImageFilter\n",
    "from matplotlib.pyplot import imshow\n",
    "from matplotlib import pyplot as plt\n",
    "import scipy.stats as st"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "convenient-tracy",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch.utils.data\n",
    "from net import *\n",
    "from model import Model\n",
    "from launcher import run\n",
    "from checkpointer import Checkpointer\n",
    "from dlutils.pytorch import count_parameters\n",
    "from defaults import get_cfg_defaults\n",
    "import lreq\n",
    "import logging\n",
    "from torch.autograd import Variable"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "tracked-ferry",
   "metadata": {},
   "outputs": [],
   "source": [
    "from ipywidgets import interact, interactive, fixed, interact_manual\n",
    "import ipywidgets as widgets\n",
    "import face_recognition\n",
    "from tqdm import trange\n",
    "import PIL\n",
    "from ffhq_dataset.landmarks_detector import LandmarksDetector\n",
    "import bz2\n",
    "from keras.utils import get_file\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "exotic-liberal",
   "metadata": {},
   "outputs": [],
   "source": [
    "path = \"raw_images/\"\n",
    "valid_images = (\".jpg\",\".png\")\n",
    "for f in os.listdir(path):\n",
    "    if f.endswith(valid_images):\n",
    "        img = PIL.Image.open(path + f)\n",
    "        display(img.resize((int(img.size[0]/3), int(img.size[1]/3)), PIL.Image.LANCZOS))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "occupational-plumbing",
   "metadata": {},
   "outputs": [],
   "source": [
    "!python align_images.py raw_images/ aligned_images/ alignement_vector/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "electric-murder",
   "metadata": {},
   "outputs": [],
   "source": [
    "path = \"aligned_images/\"\n",
    "for f in os.listdir(path):\n",
    "    if f.endswith(valid_images):\n",
    "        img = PIL.Image.open(path + f)\n",
    "        display(img.resize((256,256), PIL.Image.LANCZOS))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "color-backup",
   "metadata": {},
   "outputs": [],
   "source": [
    "cfg = get_cfg_defaults()\n",
    "#cfg.OUTPUT_DIR = \"training_artifacts/ffhq\"\n",
    "cfg.merge_from_file(\"configs/ffhq.yaml\")\n",
    "\n",
    "logger = logging.getLogger(\"logger\")\n",
    "logger.setLevel(logging.DEBUG)\n",
    "\n",
    "model = Model(\n",
    "    startf=cfg.MODEL.START_CHANNEL_COUNT,\n",
    "    layer_count=cfg.MODEL.LAYER_COUNT,\n",
    "    maxf=cfg.MODEL.MAX_CHANNEL_COUNT,\n",
    "    latent_size=cfg.MODEL.LATENT_SPACE_SIZE,\n",
    "    truncation_psi=cfg.MODEL.TRUNCATIOM_PSI,\n",
    "    truncation_cutoff=cfg.MODEL.TRUNCATIOM_CUTOFF,\n",
    "    mapping_layers=cfg.MODEL.MAPPING_LAYERS,\n",
    "    channels=cfg.MODEL.CHANNELS,\n",
    "    generator=cfg.MODEL.GENERATOR,\n",
    "    encoder=cfg.MODEL.ENCODER)\n",
    "\n",
    "#model = torch.load(\"training_artifacts/ffhq/model_submitted.pth\")\n",
    "#checkpoint = torch.load(\"training_artifacts/ffhq/model_submitted.pth\")\n",
    "\n",
    "#print(checkpoint[\"models\"].keys())\n",
    "#model.load_state_dict(checkpoint[\"models\"])\n",
    "\n",
    "model.eval()\n",
    "model.requires_grad_(False)\n",
    "\n",
    "\n",
    "decoder = model.decoder\n",
    "encoder = model.encoder\n",
    "mapping_tl = model.mapping_d\n",
    "mapping_fl = model.mapping_f\n",
    "dlatent_avg = model.dlatent_avg\n",
    "\n",
    "logger.info(\"Trainable parameters generator:\")\n",
    "count_parameters(decoder)\n",
    "\n",
    "logger.info(\"Trainable parameters discriminator:\")\n",
    "count_parameters(encoder)\n",
    "\n",
    "arguments = dict()\n",
    "arguments[\"iteration\"] = 0\n",
    "\n",
    "model_dict = {\n",
    "    'discriminator_s': encoder,\n",
    "    'generator_s': decoder,\n",
    "    'mapping_tl_s': mapping_tl,\n",
    "    'mapping_fl_s': mapping_fl,\n",
    "    'dlatent_avg': dlatent_avg\n",
    "}\n",
    "\n",
    "checkpointer = Checkpointer(cfg,\n",
    "                            model_dict,\n",
    "                            {},\n",
    "                            logger=logger,\n",
    "                            save=False)\n",
    "\n",
    "extra_checkpoint_data = checkpointer.load()\n",
    "\n",
    "model.eval()\n",
    "\n",
    "layer_count = cfg.MODEL.LAYER_COUNT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "smart-problem",
   "metadata": {},
   "outputs": [],
   "source": [
    "def encode(x):\n",
    "    Z, _ = model.encode(x, layer_count - 1, 1)\n",
    "    Z = Z.repeat(1, model.mapping_f.num_layers, 1)\n",
    "    return Z\n",
    "\n",
    "def decode(x):\n",
    "    return model.decoder(x, layer_count - 1, 1, noise=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "continental-insulin",
   "metadata": {},
   "outputs": [],
   "source": [
    "def load(path):\n",
    "    img = np.asarray(Image.open(path))\n",
    "\n",
    "    if img.shape[2] == 4:\n",
    "        img = img[:, :, :3]\n",
    "    im = img.transpose((2, 0, 1))\n",
    "    x = torch.tensor(np.asarray(im, dtype=np.float32), requires_grad=True) / 127.5 - 1.\n",
    "    if x.shape[0] == 4:\n",
    "        x = x[:3]\n",
    "\n",
    "    needed_resolution = model.decoder.layer_to_resolution[-1]\n",
    "    while x.shape[2] > needed_resolution:\n",
    "        x = F.avg_pool2d(x, 2, 2)\n",
    "    if x.shape[2] != needed_resolution:\n",
    "        x = F.adaptive_avg_pool2d(x, (needed_resolution, needed_resolution))\n",
    "\n",
    "    img_src = ((x * 0.5 + 0.5) * 255).type(torch.long).clamp(0, 255).cpu().type(torch.uint8).transpose(0, 2).transpose(0, 1).numpy()\n",
    "\n",
    "    latents_original = encode(x[None, ...])\n",
    "    latents = latents_original[0, 0].clone()\n",
    "    latents -= model.dlatent_avg.buff.data[0]\n",
    "\n",
    "    return latents, latents_original, img_src"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "chief-fisher",
   "metadata": {},
   "outputs": [],
   "source": [
    "def convert_to_image(latent_vector):\n",
    "    x_rec = decode(latent_vector.detach())\n",
    "    resultsample = ((x_rec * 0.5 + 0.5) * 255).type(torch.long).clamp(0, 255)\n",
    "    resultsample = resultsample.cpu()[0, :, :, :]\n",
    "    \n",
    "    return (resultsample.type(torch.uint8).transpose(0, 2).transpose(0, 1))\n",
    "\n",
    "\n",
    "def gkern(kernlen=21, nsig=3):\n",
    "    \"\"\"Returns a 2D Gaussian kernel.\"\"\"\n",
    "\n",
    "    x = np.linspace(-nsig, nsig, kernlen+1)\n",
    "    kern1d = np.diff(st.norm.cdf(x))\n",
    "    kern2d = np.outer(kern1d, kern1d)\n",
    "    return kern2d/kern2d.sum()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "residential-reaction",
   "metadata": {},
   "outputs": [],
   "source": [
    "image_path = \"aligned_images/IMG_5551_02.png\"\n",
    "\n",
    "latents, latents_original, img_src = load(image_path)\n",
    "\n",
    "img = Image.fromarray(convert_to_image(latents_original).detach().numpy())\n",
    "display(img.resize((256,256), PIL.Image.LANCZOS))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "steady-cartridge",
   "metadata": {},
   "outputs": [],
   "source": [
    "input_latent = Variable(latents_original.data, requires_grad=True)\n",
    "\n",
    "if img_src.shape[2] == 4:\n",
    "    img_src = img_src[:, :, :3]\n",
    "img = img_src.transpose((2, 0, 1))\n",
    "x = torch.tensor(np.asarray(img, dtype=np.float32), requires_grad=True) / 127.5 - 1.\n",
    "if x.shape[0] == 4:\n",
    "    x = x[:3]\n",
    "x = x.reshape(1,3,1024,1024)\n",
    "x = Variable(x, requires_grad=False)\n",
    "\n",
    "t = trange(10, desc='Loss', leave=True)\n",
    "for _ in t:\n",
    "    # Foreward pass\n",
    "    prediction = model.decoder(input_latent, layer_count - 1, 1, noise=False)\n",
    "\n",
    "    loss = ((prediction - x).pow(2) * torch.tensor(gkern(1024, 2) * 5000)).sum()\n",
    "    \n",
    "    # Backward pass\n",
    "    loss.backward()\n",
    "    \n",
    "    t.set_description(\"Loss: \" + str(loss.detach().numpy()))\n",
    "\n",
    "    # Update\n",
    "    input_latent = input_latent - 0.004 * input_latent.grad.data\n",
    "    input_latent = Variable(input_latent.data, requires_grad=True)\n",
    "\n",
    "img = Image.fromarray(convert_to_image(input_latent).detach().numpy())\n",
    "display(img.resize((256,256), PIL.Image.LANCZOS))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "wrong-difference",
   "metadata": {},
   "outputs": [],
   "source": [
    "face_landmarks = \"alignement_vector/IMG_5551_02.npy\"\n",
    "src_file = \"raw_images/IMG_5551.jpg\"\n",
    "mask_file = \"masks/IMG_5551_02.png\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "balanced-memorial",
   "metadata": {},
   "outputs": [],
   "source": [
    "latents_x = input_latent.detach().clone()\n",
    "\n",
    "indices = [0, 1, 2, 3, 4, 10, 11, 17, 19]\n",
    "W = [torch.tensor(np.load(\"principal_directions/direction_%d.npy\" % i), dtype=torch.float32) for i in indices]\n",
    "\n",
    "def alter(gender, smile, attract, hair, age, lips, nose, fat, glasses):\n",
    "    altered = input_latent.detach().clone()\n",
    "    \n",
    "    altered += gender *  W[0]\n",
    "    altered += smile *   W[1]\n",
    "    altered += attract * W[2]\n",
    "    altered += hair *    W[3]\n",
    "    altered -= age *     W[4]\n",
    "    altered += lips *    W[5]\n",
    "    altered += nose *    W[6]\n",
    "    altered += fat *     W[7]\n",
    "    altered += glasses * W[8]\n",
    "    \n",
    "    im = convert_to_image(altered)\n",
    "    \n",
    "    img = Image.fromarray(im.detach().numpy())\n",
    "    \n",
    "    generated_image_landmarks = face_recognition.face_landmarks(np.array(img))[0]\n",
    "        \n",
    "    lm = np.load(face_landmarks)\n",
    "    lm_chin = lm[0:17]  # left-right\n",
    "    lm_eyebrow_left = lm[17:22]  # left-right\n",
    "    lm_eyebrow_right = lm[22:27]  # left-right\n",
    "    lm_nose = lm[27:31]  # top-down\n",
    "    lm_nostrils = lm[31:36]  # top-down\n",
    "    lm_eye_left = lm[36:42]  # left-clockwise\n",
    "    lm_eye_right = lm[42:48]  # left-clockwise\n",
    "    lm_mouth_outer = lm[48:60]  # left-clockwise\n",
    "    lm_mouth_inner = lm[60:68]  # left-clockwise\n",
    "\n",
    "    # Calculate auxiliary vectors.\n",
    "    eye_left = np.mean(lm_eye_left, axis=0)\n",
    "    eye_right = np.mean(lm_eye_right, axis=0)\n",
    "    eye_avg = (eye_left + eye_right) * 0.5\n",
    "    eye_to_eye = eye_right - eye_left\n",
    "    mouth_left = lm_mouth_outer[0]\n",
    "    mouth_right = lm_mouth_outer[6]\n",
    "    mouth_avg = (mouth_left + mouth_right) * 0.5\n",
    "    eye_to_mouth = mouth_avg - eye_avg\n",
    "\n",
    "    left_eye_before = np.mean(generated_image_landmarks[\"left_eye\"], axis=0)\n",
    "    right_eye_before = np.mean(generated_image_landmarks[\"right_eye\"], axis=0)\n",
    "\n",
    "    x1b, y1b = left_eye_before\n",
    "    x2b, y2b = right_eye_before\n",
    "\n",
    "    y1a, x1a = eye_left\n",
    "    y2a, x2a = eye_right\n",
    "\n",
    "    # Generate linear system to compute transformation\n",
    "    a = np.matrix([[-y1b, x1b, 0, 1], [x1b, y1b, 1, 0], [-y2b, x2b, 0, 1], [x2b, y2b, 1, 0]])\n",
    "\n",
    "    b = np.array([y1a, x1a, y2a, x2a])\n",
    "\n",
    "    # Solve linear system\n",
    "    x = np.linalg.solve(a, b)\n",
    "\n",
    "    # Calc transformation matrix\n",
    "    transform_matrix = np.matrix([[x[1], -x[0], x[3]], [x[0], x[1], x[2]], [0, 0, 1]])\n",
    "\n",
    "    transform_matrix = np.array([[x[1], -x[0], x[3]], [x[0], x[1], x[2]], [0, 0, 1]])\n",
    "\n",
    "    transform_inv = np.linalg.inv(transform_matrix)\n",
    "\n",
    "    background = Image.open(src_file).convert(\"RGB\")\n",
    "\n",
    "    foreground = img\n",
    "    mask = Image.open(mask_file).convert(\"L\").resize(foreground.size).filter(ImageFilter.GaussianBlur(20))\n",
    "    \n",
    "    foreground = foreground.transform(background.size, PIL.Image.AFFINE, transform_inv.flatten()[:6], resample=Image.NEAREST)    \n",
    "    mask = mask.transform(background.size, PIL.Image.AFFINE, transform_inv.flatten()[:6], resample=Image.NEAREST)\n",
    "\n",
    "    \n",
    "    background = Image.composite(foreground, background, mask)\n",
    "    \n",
    "    size = 512, 512\n",
    "    background.thumbnail(size, Image.ANTIALIAS)\n",
    "    \n",
    "    display(background)\n",
    "\n",
    "    \n",
    "w = interact(\n",
    "    alter, \n",
    "    gender=(-19,20,1), \n",
    "    smile=(-9,10,1), \n",
    "    attract=(-9,10,1),\n",
    "    hair=(-9,10,1),\n",
    "    age=(-9,10,1), \n",
    "    lips=(-9,10,1),\n",
    "    nose=(-9,10,1),\n",
    "    fat=(-9,10,1),\n",
    "    glasses=(-99,100,1)\n",
    ")\n",
    "display(w)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "white-integrity",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
